mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
chore: merge litellm_internal_staging into fix/batch-retrieve-model-group
This commit is contained in:
commit
f66663cc8b
104 changed files with 5246 additions and 790 deletions
38
.github/e2e-stack/assert_tests_ran.py
vendored
Normal file
38
.github/e2e-stack/assert_tests_ran.py
vendored
Normal 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
17
.github/e2e-stack/down.sh
vendored
Executable 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
49
.github/e2e-stack/secrets_to_env.py
vendored
Normal 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
44
.github/e2e-stack/select_tests.py
vendored
Normal 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
207
.github/e2e-stack/up.sh
vendored
Executable 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"
|
||||
21
.github/workflows/_test-unit-base.yml
vendored
21
.github/workflows/_test-unit-base.yml
vendored
|
|
@ -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 }}
|
||||
|
|
|
|||
3
.github/workflows/test-code-quality.yml
vendored
3
.github/workflows/test-code-quality.yml
vendored
|
|
@ -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
240
.github/workflows/test-e2e-changed.yml
vendored
Normal 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
|
||||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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 \
|
||||
|
|
|
|||
|
|
@ -0,0 +1,2 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_MCPServerTable" ADD COLUMN IF NOT EXISTS "per_server_oauth_discovery" BOOLEAN NOT NULL DEFAULT false;
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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"],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}),
|
||||
)
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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]
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ------------
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
157
litellm/router_utils/auto_router_tuning_baseline.py
Normal file
157
litellm/router_utils/auto_router_tuning_baseline.py
Normal 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)
|
||||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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.",
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
203
tests/code_coverage_tests/test_e2e_changed_gate.py
Normal file
203
tests/code_coverage_tests/test_e2e_changed_gate.py
Normal 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()
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,),
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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],
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
64
tests/e2e/gateway/stage_mirror_ci_config.yml
Normal file
64
tests/e2e/gateway/stage_mirror_ci_config.yml
Normal 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
|
||||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
87
tests/e2e/test_proxy_client.py
Normal file
87
tests/e2e/test_proxy_client.py
Normal 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",)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
[
|
||||
|
|
|
|||
|
|
@ -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."""
|
||||
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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"] == [
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
)
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
59
tests/test_litellm/router_strategy/test_lowest_cost.py
Normal file
59
tests/test_litellm/router_strategy/test_lowest_cost.py
Normal 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}
|
||||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
@ -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",
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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>)
|
||||
|
|
|
|||
|
|
@ -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 }
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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");
|
||||
});
|
||||
});
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
);
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 = {
|
||||
|
|
|
|||
|
|
@ -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";
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
};
|
||||
|
|
@ -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
Loading…
Add table
Reference in a new issue